#include <mex.h>
#include <math.h>
#include "shapeAndResponse.h"
#include  <omp.h>


void cpVectorDouble2Float(float * p_dest, double *p_src, int num_elem){
	for( int i = 0;i<num_elem; ++i)
		p_dest[i]=(float)p_src[i];
}


void *getPtrToField(const mxArray *p_struct, int i, const char* field_name){
	const mxArray *p_field=mxGetField(p_struct,i,field_name);
	return (void*)mxGetPr(p_field);
}

void mexFunction( int nlhs, mxArray *plhs[],
				 int nrhs, const mxArray *prhs[] ){
	// access the variables
	
	int i=0;
	int n_image = floor(mxGetScalar(prhs[i++])+0.5);// n_image 2
	int n_orient = floor(mxGetScalar(prhs[i++])+0.5);// n_orient 3
	int n_length = floor(mxGetScalar(prhs[i++])+0.5);// n_length 5
	
	int h =(int) floor(mxGetScalar(prhs[i++])+0.5f);
	// image sizes and zoom ratios
	double *p_sx = mxGetPr(prhs[i++]);
	double *p_sy = mxGetPr(prhs[i++]);
	
	double *p_zx = mxGetPr(prhs[i++]);
	double *p_zy = mxGetPr(prhs[i++]);
	
	// get M2 maps
	float **p_M2 = (float**)mxCalloc((int)(n_image*n_orient*n_length),sizeof(float*));
	for ( int i_page = 0; i_page<n_image*n_orient*n_length; ++i_page){
		mxArray *p_page = mxGetCell(prhs[i],i_page);
		p_M2[i_page]=(float*)mxGetPr(p_page);
	}
	i++;
	
	float *p_PMat = (float*) mxGetPr(prhs[i++]);
	float *p_viewMat = (float*)mxGetPr(prhs[i++]);
	
	
	// get ptr to AoT structure
	const mxArray *p_tan = prhs[i++];
	if( !mxIsStruct(p_tan)){
		mexPrintf("Error, not a structure, p_tan.\n");
	}
	
	double *p_root_tan_rotation = (double*)getPtrToField(p_tan, 0, "rotationMat");
	float root_tan_rotation[9];
	cpVectorDouble2Float(root_tan_rotation,p_root_tan_rotation,9);
	
	double *p_root_tan_origin =(double*) getPtrToField(p_tan, 0, "origin");
	
	float root_tan_T[3];
	// note that according to notes, T is the negative of the tan origin.
	for(int j =0;j<3; ++j) root_tan_T[j] = - p_root_tan_origin[j];
	
	
	
	double root_tan_scale = *(double*)getPtrToField(p_tan, 0, "swt"); //
	
	
	
	const mxArray *p_all_term_layer=mxGetField(p_tan,0,"termLayer");
	int n_layer = mxGetM(p_all_term_layer);
	const int *p_n_layer = &n_layer;
	plhs[0] = mxCreateCellArray(1,p_n_layer);
	plhs[1] = mxCreateCellArray(1,p_n_layer);

	printf("totally %d layers\n",n_layer);	
	for(int i_layer = 0; i_layer<n_layer; ++i_layer){
		printf(" To %d of %d layers\n",i_layer,n_layer);
		const mxArray *p_term_layer = mxGetCell(p_all_term_layer,i_layer);
		
		// begin parallel area;
		int n_node = mxGetN(p_term_layer);// n_node
		const int* p_n_node=&n_node;
		mxArray *score_map_cell = mxCreateCellArray(1,p_n_node);
		mxArray *arg_max_cell = mxCreateCellArray(1,p_n_node);
		
		float ***p_my_score_map_cell = (float***)malloc(n_node*sizeof(float**));
		int ***p_my_arg_max_cell = (int***)malloc(n_node*sizeof(int**));
		int *p_n_shape = (int*)malloc(n_node*sizeof(int));
		
		printf("n_node: %d\n",n_node);
		omp_set_num_threads(8);
#pragma omp parallel for 
		for( int i_node =0; i_node<n_node; ++i_node){
			
			//printf("node %d of %d\n",i_node,n_node);
			// get ptr to the plane
			const mxArray *p_plane = mxGetField(p_term_layer,i_node,"plane");
			// get plane parameters, and convert to internal representation
			plane target_plane;
			
			double *p_tmp_array=(double*)getPtrToField(p_plane,0,"nt");
			cpVectorDouble2Float(&target_plane.nt[0],p_tmp_array,4);
			p_tmp_array = (double*)getPtrToField(p_plane, 0, "origin");
			cpVectorDouble2Float(&target_plane.origin[0],p_tmp_array,3);
			p_tmp_array = (double*)getPtrToField(p_plane, 0, "u");
			cpVectorDouble2Float(&target_plane.bu[0],p_tmp_array,3);
			p_tmp_array = (double*)getPtrToField(p_plane,0,"v");
			cpVectorDouble2Float(&target_plane.bv[0],p_tmp_array,3);
			target_plane.width = *(double*)getPtrToField(p_plane,0,"w");
			target_plane.height = *(double*)getPtrToField(p_plane,0,"h");
			target_plane.stp =*(double*)getPtrToField(p_plane,0,"stp");
			
			
			target_plane.height*=target_plane.stp; // plane size in plane coordinate system
			target_plane.width *= target_plane.stp;
			
			double dims[6];
			dims[2]=n_image;
			dims[3]=n_orient;
			dims[5]=n_length;
			
			dims[0]=target_plane.height;
			dims[1]=target_plane.width;
			
			// begin the main computation part
			shapeAndResponse calculator;
			calculator.initial(dims,h);
			calculator.my_projector.initialize(root_tan_scale,root_tan_rotation,root_tan_T,(int)dims[3],(int)dims[5],h);
			calculator.my_projector.updatePlane(target_plane.bu,target_plane.bv,target_plane.origin,
												target_plane.nt,target_plane.stp,target_plane.width,target_plane.height);
			calculator.generateShapeList();
			int n_shape = calculator.numShapes();
			p_n_shape[i_node]=n_shape;
			// acllocate space for max scores of all shapes
			float **score_map = (float**)malloc(sizeof(float*)*n_shape);
			int **arg_max_map = (int**)malloc(sizeof(int*)*n_shape);
			p_my_score_map_cell[i_node]=score_map;
			p_my_arg_max_cell[i_node]=arg_max_map;
			
			for(int i_shape =0; i_shape<n_shape; ++i_shape){
				score_map[i_shape]=(float*)malloc(sizeof(float)*dims[2]);
				arg_max_map[i_shape]=(int*)malloc(sizeof(int)*dims[2]);
				
				for( int i = 0; i<dims[2]; ++i){
					score_map[i_shape][i]=NEGMAX;
					arg_max_map[i_shape][i]=-1e8;
				}
			}
			
			calculator.computeShapeScore(p_M2,p_PMat,p_viewMat,p_sx,p_sy,p_zx,p_zy, score_map,arg_max_map);
			
			// free memory
			calculator.release();
		}// end of paralle for loop
		
		for( int i_node=0; i_node<n_node; ++i_node){
			// find fields in structure
			float **score_map = p_my_score_map_cell[i_node];
			int **arg_max_map = p_my_arg_max_cell[i_node];
			int n_shape=p_n_shape[i_node];		
			
			mxArray *p_score_map;
			p_score_map=mxCreateNumericMatrix(n_image,n_shape,mxSINGLE_CLASS,mxREAL);
			
			float* p_float_array=(float*)mxGetPr(p_score_map);
			for(int i = 0; i<n_shape;++i){
				for(int j =0; j<n_image;++j){
					p_float_array[i*(int)(n_image)+j]=score_map[i][j];
				}
			}
			
			mxSetCell(score_map_cell,i_node,p_score_map);
			
			mxArray *p_arg_max_map;
			
			p_arg_max_map=mxCreateNumericMatrix(n_image,n_shape,mxINT32_CLASS,mxREAL);
			
			int *p_int_array=(int*)mxGetPr(p_arg_max_map);
			for(int i =0; i<n_shape;++i){
				for(int j =0; j<n_image;++j){
					p_int_array[i*(int)(n_image)+j]=arg_max_map[i][j];
				}
			}
			
			mxSetCell(arg_max_cell,i_node,p_arg_max_map);			
			
			
			// free matrix
			for(int i_shape =0; i_shape<n_shape; ++i_shape){
				free((void*)score_map[i_shape]);
				free((void*)arg_max_map[i_shape]);
			}
			free((void*)score_map);
			free((void*)arg_max_map);
		}// i_node
		
		free(p_n_shape);
		free(p_my_score_map_cell);
		free(p_my_arg_max_cell);
		
		mxSetCell(plhs[0],i_layer,score_map_cell);
		mxSetCell(plhs[1],i_layer,arg_max_cell);
	}// i_layer
	
}
